package org.example

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}

object sparkSQL {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .master("local[*]")
      .appName("sparkBase")
      .getOrCreate()

    import spark.implicits._

    // 定义用户表结构
    val schemaUser = StructType(Seq(
      StructField("id", IntegerType),
      StructField("gender", StringType),
      StructField("age", IntegerType),
      StructField("occupation", IntegerType),
      StructField("location", StringType)
    ))

    // 读取用户表
    val users = spark.read.option("sep", "::").schema(schemaUser).csv("src/main/resources/users.dat")

    // 定义评分表结构
    val schemaRating = StructType(Seq(
      StructField("user_id", IntegerType),
      StructField("movie_id", IntegerType),
      StructField("rating", IntegerType),
      StructField("timestamp", IntegerType)
    ))

    // 读取评分表
    val ratings = spark.read.option("sep", "::").schema(schemaRating).csv("src/main/resources/ratings.dat")

    // 定义电影表结构
    val schemaMovie = StructType(Seq(
      StructField("movie_id", IntegerType),
      StructField("title", StringType),
      StructField("genres", StringType)
    ))

    // 读取电影表
    val movies = spark.read.option("sep", "::").schema(schemaMovie).csv("src/main/resources/movies.dat")

    // 查询18岁女生评分为5分的所有电影名称
    val result = users
      .filter($"gender" === "F" && $"age" === 18)
      .join(ratings, users("id") === ratings("user_id"))
      .filter($"rating" === 5)
      .join(movies, ratings("movie_id") === movies("movie_id"))
      .select("title")
      .distinct()

    result.show()

    spark.stop()
  }
}